# from sqlalchemy import create_engine, text


# DB_CONFIG = {
#     'host': 'localhost',
#     'port': '5432',
#     'database': 'Policy_DB',
#     'user': 'postgres',
#     'password': 'postgres'
# }

# db_url = f"postgresql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
# engine = create_engine(db_url)

# try:
#     with engine.connect() as conn:
#         gb_tag_ids_res = conn.execute(text("SELECT gb_industry_2 FROM gb_industry_2_labels_table"))
#         gb_tag_ids = gb_tag_ids_res.fetchall()
#         for gb_id in gb_tag_ids:
#             print(type(gb_id))
#             break
# except Exception as e:
#     print(e)



from typing import List
import psycopg2
from psycopg2 import sql

def query_gb_tags(tag_names, table_name):
    """
    从GB标签表查询标签序列并打印结果
    参数:
        tag_names (list): 要查询的标签名列表，例如 ['农业', '林业']
        table_name(str): 要查询的表，'gb_industry_2_labels_table'或者'zx_industry_2_labels_table'
    """
    # 数据库连接配置
    conn = psycopg2.connect(
        dbname="Policy_DB",  # 替换为您的数据库名
        user="postgres",     # 替换为您的数据库用户名
        password="postgres", # 替换为您的数据库密码
        host="localhost"     # 替换为您的数据库主机
    )
    
    try:
        with conn.cursor() as cursor:
            if table_name == 'gb_industry_2_labels_table':
                sql_execute = """
                    SELECT gb_industry_2, gb_industry_2_labels_table.tag_name 
                    FROM gb_industry_2_labels_table
                    WHERE gb_industry_2_labels_table.tag_name = ANY(%s)
                    ORDER BY gb_industry_2_labels_table.tag_name
                """ 
            else:
                sql_execute = """
                    SELECT zx_policy_2, zx_industry_2_labels_table.tag_name 
                    FROM zx_industry_2_labels_table
                    WHERE zx_industry_2_labels_table.tag_name = ANY(%s)
                    ORDER BY zx_industry_2_labels_table.tag_name
                """
            # 构建SQL查询
            query = sql.SQL(sql_execute)
            
            # 执行查询
            cursor.execute(query, (tag_names,))
            
            # 获取结果
            results = cursor.fetchall()
            
            tags = []
            # 打印查询结果
            for tag_code, tag_name in results:
                # print(f"{tag_code:<15}{tag_name:<25}")
                tags.append(str(tag_code))
            
            print(tags)
            res = find_policies_by_tags(tags, table_name)
            # for id in res:
            #     print(id)
            return res

    except Exception as e:
        print(f"查询出错: {e}")
        return []
    finally:
        conn.close()

def find_policies_by_tags(tag_codes: List[str], table_name: str) -> List[str]:
    """
    查找同时满足所有指定GB二级标签的政策ID
    
    参数:
        tag_codes: 需要匹配的GB二级标签代码列表(如['1-1', '4-3']
        table_name: 查询的表
    返回:
        满足所有标签条件的政策ID列表
    """
    conn = None
    try:
        # 连接数据库
        conn = psycopg2.connect(
            dbname="Policy_DB",
            user="postgres",
            password="postgres",
            host="localhost"
        )
        
        with conn.cursor() as cursor:
            # 构建SQL查询
            if table_name == "gb_industry_2_labels_table":
                query = """
                    SELECT policy_id
                    FROM policy_gb_label_table
                    WHERE gb_industry_2 = ANY(%s)
                    GROUP BY policy_id
                    HAVING COUNT(DISTINCT gb_industry_2) = %s
                """
            else:
                query = """
                    SELECT policy_id
                    FROM policy_zx_label_table
                    WHERE zx_industry_2 = ANY(%s)
                    GROUP BY policy_id
                    HAVING COUNT(DISTINCT zx_industry_2) = %s
                """
            
            # 执行查询
            cursor.execute(query, (tag_codes, len(tag_codes)))
            
            # 获取结果
            results = [row[0] for row in cursor.fetchall()]
            return results
            
    except Exception as e:
        print(f"数据库查询出错: {e}")
        return []
    finally:
        if conn:
            conn.close()

def get_policy_content_by_id(policy_ids):
    conn = None
    try:
        # 连接数据库
        conn = psycopg2.connect(
            dbname="Policy_DB",
            user="postgres",
            password="postgres",
            host="localhost"
        )
        
        with conn.cursor() as cursor:
            # 构建SQL查询
            query = """
                SELECT *
                FROM policy_content_table
                WHERE policy_id = ANY(%s)
            """
            
            # 执行查询
            cursor.execute(query, (policy_ids,))
            
            # 获取结果
            results = cursor.fetchall()
            return results
            
    except Exception as e:
        print(f"数据库查询出错: {e}")
        return []
    finally:
        if conn:
            conn.close()

def search_policy_by_tags(tags: list):
    # 调用查询函数
    res1 = query_gb_tags(tags, "gb_industry_2_labels_table")
    res2 = query_gb_tags(tags, "zx_industry_2_labels_table")

    if len(res1) != 0 and len(res2) != 0:
        res_policy_ids = list(set(res1) & set(res2))
    elif len(res1) == 0:
        res_policy_ids = list(set(res2))
    elif len(res2) == 0:
        res_policy_ids = list(set(res1))

    result = get_policy_content_by_id(res_policy_ids)
    
    final_res = []
    keys = ("policy_id", "title", "write_time", "publish_time", "start_time" ,"end_time", "validity", "file_type", "department","document_number",
    "main_text",
    "is_company_policy", 
    "is_crossborder_policy",
    "level",
    "type",
    "language",
    "address_1",
    "address_2",
    "address_3",
    "address_4",
    "address_5",
    "address_6")
    for r in result:
        dict = {keys[i]: r[i] for i in range(len(keys))}
        final_res.append(dict)
    
    return final_res

# 使用示例
if __name__ == "__main__":
    # 要查询的标签列表
    tags_to_query = ['农业', '林业']
    
    # 调用查询函数
    res1 = query_gb_tags(tags_to_query, "gb_industry_2_labels_table")
    res2 = query_gb_tags(tags_to_query, "zx_industry_2_labels_table")

    if len(res1) != 0 and len(res2) != 0:
        res_policy_ids = list(set(res1) & set(res2))
    elif len(res1) == 0:
        res_policy_ids = list(set(res2))
    elif len(res2) == 0:
        res_policy_ids = list(set(res1))

    # print(res_policy_ids)
    result = get_policy_content_by_id(res_policy_ids)
    
    final_res = []
    # print(type(result))
    keys = ("policy_id", "title", "write_time", "publish_time", "start_time" ,"end_time", "validity", "file_type", "department","document_number",
    "main_text",
    "is_company_policy", 
    "is_crossborder_policy",
    "level",
    "type",
    "language",
    "address_1",
    "address_2",
    "address_3",
    "address_4",
    "address_5",
    "address_6")
    for r in result:
        dict = {keys[i]: r[i] for i in range(len(keys))}
        final_res.append(dict)

    print(len(final_res))


    # ['1-1', '1-2']
    # ['1-4', '6-3', '1-5', '4-3', '2-2', '7-3']
    # 要查询的GB二级标签代码列表
    # required_tags = ['1-1', '4-3']
    
    # 查询满足所有标签的政策
    # matching_policies = find_policies_by_gb_tags(required_tags)
    
    # 打印结果
    # print(f"同时满足标签 {required_tags} 的政策ID:")
    # for policy_id in matching_policies:
    #     print(policy_id)